#include "kernel_operator.h"
using namespace AscendC;

class KernelRedcuesum
{

public:
    __aicore__ inline KernelRedcuesum() {}

    __aicore__ inline void Init(GM_ADDR x, GM_ADDR axis, GM_ADDR y, int32_t dim_num, int64_t shape[],//这里应该是dim_num还是dimNum,这里dim_num是shape的长度，比如（2，3，4）长度为3 
                                uint32_t blockLength, uint32_t axLength)
    {
            xGm.SetGlobalBuffer((__gm__ DTYPE_X *)x, blockLength);
            axisGm.SetGlobalBuffer((__gm__ DTYPE_AXIS *)axis, axLength);
            yGm.SetGlobalBuffer((__gm__ DTYPE_Y *)y, blockLength);
            int32_t axisCount = 0; // 轴的数量
            int32_t* axes = NULL; // 轴的数组
            int32_t axisScalar; // 单个轴索引
            bool isSingleAxis = true; // 是否只有一个轴
            bool isAllAxes = false; // 是否对所有轴求和

            if (axis == NULL or axLength==dim_num){
                isAllAxes=true;
            }else{
                axisCount = axLength;
                axes = (int32_t*)malloc(axisCount * sizeof(int32_t));
                for (int i = 0; i < axisCount; i++) {
                    axes[i] = axisGm.GetValue(i);
                    if (axes[i] < 0) {
                        axes[i] += dim_num; // 转换负数索引
                    }
            }
                if (axisCount == 1) {
                    isSingleAxis = true;
                    axisScalar = axes[0];
                } else {
                    isSingleAxis = false;
                    //对axes进行排序
                    for (int i = 0;i<axLength;i++){
                        for (int j = i; j<axLength;j++){
                            if (axes[j]>axes[j+1]){
                                temp=axes[j];
                                axes[j] = axes[j+1];
                                axes[j+1] = temp
                            }
                            
                            
                        }
                    }
                }
            }
        if (isAllAxes) {
            // 对所有元素求和
            float sum = 0;
            for (int i = 0; i < blockLength; i++) {
                sum += (float)xGm.GetValue(i);
            }
            yGm.SetValue(0, (DTYPE_Y)sum);
        } else if (isSingleAxis) {
            if (axisScalar < 0)
            {
                axisScalar += dim_num;
            }
            int32_t f = 1;
            int32_t b = 1;
            int32_t a = shape[axisScalar];
            for (int i = 0;i<axisScalar;i++){
                f *= shape[i];
            }
            for (int j = axisScalar+1;j<dim_num;j++){
                b*= shape[j];
            }
            for (int p = 0; p < f;p++){
                for (int k=0;k<b;k++){
                    float tmpValue = 0;
                    for (int q= 0;q<a;q++){
                        tmpValue += (float) xGm.GetValue(p*b*a + k + q*b);
                    }
                    yGm.SetValue(p*b+k, (DTYPE_X)tmpValue);
                }
            }
        }else{
            // 多轴求和
            __gm__ DTYPE_X tempGm; // 临时张量
            tempGm.SetGlobalBuffer((__gm__ DTYPE_X*)temp, blockLength);
        // 处理每个轴
        for (int ax = axLength-1; ax >=0; ax--) {
            axisScalar = axes[ax];

            // 重置临时张量
            tempGm.SetGlobalBuffer((__gm__ DTYPE_X*)temp, blockLength);

            // 更新形状信息
            int32_t f = 1;
            int32_t b = 1;
            int32_t a = shape[axisScalar];
            for (int i = 0; i < axisScalar; i++) {
                f *= shape[i];
            }
            for (int j = axisScalar + 1; j < dim_num; j++) {
                b *= shape[j];
            }

            // 沿当前轴求和
            for (int p = 0; p < f; p++) {
                for (int k = 0; k < b; k++) {
                    float tmpValue = 0;
                    for (int q = 0; q < a; q++) {
                        tmpValue += (float)xGm.GetValue(p * b * a + k + q * b);
                    }
                    tempGm.SetValue(p * b + k, (DTYPE_X)tmpValue);
                }
            }

            // 更新 xGm 为新的输入
            xGm.SetGlobalBuffer(tempGm.GetBuffer(), tempGm.GetLength());
        }

        // 最终结果存放在 yGm 中
        //怎么放呀
  // 清理内存
    if (axes != NULL) {
        free(axes);
    }
    }
}
private:
    GlobalTensor<DTYPE_X> xGm;
    GlobalTensor<DTYPE_AXIS> axisGm;
    GlobalTensor<DTYPE_Y> yGm;
};
extern "C" __global__ __aicore__ void reduce_sum(GM_ADDR x, GM_ADDR axis, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling)
{
    GET_TILING_DATA(tiling_data, tiling);
    // TODO: user kernel impl
    KernelRedcuesum op;
    op.Init(x,axis,y,tiling_data.dimNum,tiling_data.shape,tiling_data.blockLength);

}
